#include "../DWMRICore/Fiber.h"
#include "../DWMRICore/Scalar.h"
#include "../DWMRICore/DTI.h"

#include "LeastSquare.h"

#include <teem\ell.h>

#include <omp.h>

#define _USE_MATH_DEFINES
#include <math.h>

char *fiber_pathname = NULL;
char *dst_pathname = NULL;

char **confPathname = NULL;
char **vertPathname = NULL;
int w, h, d, scale;
float stepSize, maxLength;

int kernel_size = 0;
float fiber_length = 0.0f;

#ifndef PARAM_COUNT
#define PARAM_COUNT 17
#endif

#ifndef RESAMPLE
#define RESAMPLE 10
#endif

//struct Fiber {
//	Seed seed;
//	CPoint3F parameters[100];
//	uint count;
//};

struct Fiber {
	Seed seed;
	CPoint3F parameters[PARAM_COUNT];
	uint count;
};

void PrintUsage()
{
	printf("Usage: ComputeMyEdgeStrength -fiber <configure_file> -dst <dst_file> -length <fiber_length> -kernel <kernel_size>\n");
}

int ParseArguments(const int argc, char *argv[])
{
	for (int i = 1; i < argc; ++i) {
		if (strcmp(argv[i], "-fiber") == 0) {
			fiber_pathname = argv[++i];
		} else if (strcmp(argv[i], "-dst") == 0) {
			dst_pathname = argv[++i];
		} else if (strcmp(argv[i], "-length") == 0) {
			fiber_length = (float)atof(argv[++i]);
		} else if (strcmp(argv[i], "-kernel") == 0) {
			kernel_size = atoi(argv[++i]);
		} else {
			return -1;
		}
	}

	if (fiber_pathname == NULL || dst_pathname == NULL || fiber_length < 0.0001f || kernel_size == 0)
		return -2;

	return 0;
}

//void FiberParameterization(CFiber *fibers, Fiber *parameters, const int count)
//{
//	uint offset = int(fiber_length / stepSize);
//	for (int i = 0; i < count; ++i) {
//		parameters[i].seed = fibers[i].m_seed;
//		uint IF = fibers[i].m_fCount > offset ? offset : fibers[i].m_fCount;
//		uint IB = fibers[i].m_bCount > offset ? offset : fibers[i].m_bCount;
//
//		int id = 0;
//		for (uint k = IB-1; k > 0; k--) {
//			parameters[i].parameters[id] = fibers[i].m_pB[k];
//			id++;
//		}
//		for (uint k = 0; k < IF; ++k) {
//			parameters[i].parameters[id] = fibers[i].m_pF[k];
//			id++;
//		}
//		parameters[i].count = id;
//	}
//}

void FiberParameterization(CFiber *fibers, Fiber *parameters, const int count)
{
	uint offset = uint(fiber_length / stepSize);
	for (int i = 0; i < count; ++i) {
		parameters[i].seed = fibers[i].m_seed;

		uint IF = fibers[i].m_fCount > offset ? offset : fibers[i].m_fCount;
		uint IB = fibers[i].m_bCount > offset ? offset : fibers[i].m_bCount;

		CPoint3F *buffer = new CPoint3F[(IF+IB-1)*RESAMPLE+1];
		buffer[0] = fibers[i].m_pB[IB-1];
		int id = 1;
		for (uint k = IB-1; k > 0; k--) {
			for (int n = 0; n < RESAMPLE; ++n) {
				float ratio = float(n+1) / (float)RESAMPLE;
				buffer[id] = fibers[i].m_pB[k-1] * ratio + fibers[i].m_pB[k] * (1.0f - ratio);
				id++;
			}
		}
		for (uint k = 0; k < IF-1; k++) {
			for (int n = 0; n < RESAMPLE; ++n) {
				float ratio = float(n+1) / (float)RESAMPLE;
				buffer[id] = fibers[i].m_pF[k+1] * ratio + fibers[i].m_pF[k] * (1.0f - ratio);
				id++;
			}
		}

		int step = (id-1) / (PARAM_COUNT-1);

		for (uint k = 0; k < PARAM_COUNT; ++k) {
			parameters[i].parameters[k] = buffer[k*step];
		}
		parameters[i].count = PARAM_COUNT;

		SafeDeleteArray(buffer);
	}
}


void ComputeTensor(Fiber *fibers, const int count, float *tensor)
{
	/* compute fourier descriptors for each fiber */
	float *fd = new float[count*30];
	for (int i = 0; i < count; ++i) {
		// centroid point
		CPoint3F cp = CPoint3F(0.0f, 0.0f, 0.0f);
		for (uint k = 0; k < fibers[i].count; ++k) {
			cp.m_x += fibers[i].parameters[k].m_x;
			cp.m_y += fibers[i].parameters[k].m_y;
			cp.m_z += fibers[i].parameters[k].m_z;
		}
		cp.m_x /= (float)(fibers[i].count);
		cp.m_y /= (float)(fibers[i].count);
		cp.m_z /= (float)(fibers[i].count);

		// fourier descriptors
		for (int l = 0; l < 30; ++l) {
			float Clx_r = 0.0f, Clx_i = 0.0f;
			float Cly_r = 0.0f, Cly_i = 0.0f;
			float Clz_r = 0.0f, Clz_i = 0.0f;

			for (uint k = 0; k < fibers[i].count; ++k) {
				float angle = 2.0f * M_PI * k * (l);
				angle = angle / (float)(fibers[i].count);

				float COS = cosf(angle);
				float SIN = sinf(angle);

				Clx_r += (fibers[i].parameters[k].m_x - cp.m_x) * COS;
				Clx_i += (fibers[i].parameters[k].m_x - cp.m_x) * SIN;

				Cly_r += (fibers[i].parameters[k].m_y - cp.m_y) * COS;
				Cly_i += (fibers[i].parameters[k].m_y - cp.m_y) * SIN;

				Clz_r += (fibers[i].parameters[k].m_z - cp.m_z) * COS;
				Clz_i += (fibers[i].parameters[k].m_z - cp.m_z) * SIN;
			}
			Clx_r /= sqrtf((float)(fibers[i].count));
			Clx_i /= sqrtf((float)(fibers[i].count));
			Cly_r /= sqrtf((float)(fibers[i].count));
			Cly_i /= sqrtf((float)(fibers[i].count));
			Clz_r /= sqrtf((float)(fibers[i].count));
			Clz_i /= sqrtf((float)(fibers[i].count));

			fd[i*30+l] = Clx_r * Clx_r + Clx_i * Clx_i + Cly_r * Cly_r + Cly_i * Cly_i + Clz_r * Clz_r + Clz_i * Clz_i;
		}
	}

	/* construct matrix A */
	double *A = new double[count*4];
	for (int i = 0; i < count; ++i) { 
		A[i*4+0] = (fibers[i].seed.pos.m_x - fibers[0].seed.pos.m_x) * (double)scale;
		A[i*4+1] = (fibers[i].seed.pos.m_y - fibers[0].seed.pos.m_y) * (double)scale;
		A[i*4+2] = (fibers[i].seed.pos.m_z - fibers[0].seed.pos.m_z) * (double)scale;
		A[i*4+3] = 1.0;
	}

	/* compute pseudoinverse of matrix A */
	double *pinvA = new double[4*count];
	PseudoInverse4(A, pinvA, count, 4);

	/* estimate the function for each parameter */
	float *x = new float[4];
	float *y = new float[count];
	for (int i = 0; i < 30; ++i) {
		for (int k = 0; k < count; ++k) { 
			y[k] = fd[k*30+i];
		}
		memset(x, 0, sizeof(float)*4);

		//LeastSquare(A, x, y, count, 3);
		LeastSquareFast4(pinvA, x, y, count, 4);

		tensor[1] += x[0] * x[0];
		tensor[2] += x[0] * x[1];
		tensor[3] += x[0] * x[2];
		tensor[4] += x[1] * x[1];
		tensor[5] += x[1] * x[2];
		tensor[6] += x[2] * x[2];
	}

	SafeDeleteArray(A);
	SafeDeleteArray(pinvA);
	SafeDeleteArray(x);
	SafeDeleteArray(y);
	SafeDeleteArray(fd);
}

void ReverseFiber(Fiber *fiber)
{
	fiber->seed.dir.m_x = -fiber->seed.dir.m_x;
	fiber->seed.dir.m_y = -fiber->seed.dir.m_y;
	fiber->seed.dir.m_z = -fiber->seed.dir.m_z;

	for (uint i = 0; i < fiber->count/2; ++i) {
		//swap i with fiber->count-i-1
		CPoint3F temp = fiber->parameters[i];
		fiber->parameters[i] = fiber->parameters[fiber->count-i-1];
		fiber->parameters[fiber->count-i-1] = temp;
	}
}

int main(int argc, char *argv[])
{
	if (ParseArguments(argc, argv) != 0) {
		PrintUsage();
		return 0;
	}

	/* read the configure file */
	ReadConfigureFile(fiber_pathname, &confPathname, &vertPathname, w, h, d, scale, stepSize, maxLength);

	float length = fiber_length > maxLength ? maxLength : fiber_length;
	int cw, ch;

	/* compute the structure tensor for each voxel */
	float *st = new float[w*h*d*7];
	memset(st, 0, sizeof(float)*w*h*d*7);

	CFiber *fibers = NULL;
	Fiber **parameters = new Fiber*[2*(kernel_size+1)+1];
	for (int i = 0; i < 2*(kernel_size+1)+1; ++i) {
		parameters[i] = new Fiber[w*h];
	}

	for (int z = 0; z < 2*(kernel_size+1); ++z) {
		ReadFibers(confPathname[z], vertPathname[z], &fibers, cw, ch);
		FiberParameterization(fibers, parameters[z], cw*ch);
		delete[] fibers;
	}
	
	int offset[9];
    offset[ 0] = 0;		offset[ 1] = -1;	offset[ 2] = 1;
    offset[ 3] = -w-1;	offset[ 4] = -w;	offset[ 5] = -w+1;
    offset[ 6] = w-1;	offset[ 7] = w;		offset[ 8] = w+1;

	int ks = 2 * kernel_size + 1;
	float *gaussian = new float[ks*ks*ks];
	for (int z = 0; z < ks; ++z) {
		for (int y = 0; y < ks; ++y) {
			for (int x = 0; x < ks; ++x) {
				int index = (z * ks + y) * ks + x;
				float dis = float((x - kernel_size) * (x - kernel_size) + (y - kernel_size) * (y - kernel_size) + (z - kernel_size) * (z - kernel_size));
				gaussian[index] = exp(-dis / (0.25f * ks * ks));
			}
		}
	}

	omp_set_num_threads(64);

	printf("\n");

	for (int z = kernel_size + 1; z < d - kernel_size - 1; ++z) {
		ReadFibers(confPathname[z+kernel_size+1], vertPathname[z+kernel_size+1], &fibers, cw, ch);
		FiberParameterization(fibers, parameters[2*(kernel_size+1)], cw*ch);
		delete[] fibers;

		for (int y = kernel_size + 1; y < h - kernel_size - 1; ++y) {
#pragma omp parallel for            
			for (int x = kernel_size + 1; x < w - kernel_size - 1; ++x) {
				int index = y * w + x;
				CVector3F base_dir = parameters[kernel_size+1][index].seed.dir;

				for (int cz = -kernel_size; cz <= kernel_size; ++cz) {
					for (int cy = -kernel_size; cy <= kernel_size; ++cy) { 
						for (int cx = -kernel_size; cx <= kernel_size; ++cx) {
							int temp_index = (y + cy) * w + (x + cx);

							Fiber f[27];

							for (int i = 0; i < 9; ++i) {
								memcpy(&(f[i]), &(parameters[1+cz+kernel_size][temp_index+offset[i]]), sizeof(Fiber));
								if (InnerProduct(base_dir, f[i].seed.dir) < 0.0f) {
									ReverseFiber(&(f[i]));
								}
								memcpy(&(f[i+9]), &(parameters[0+cz+kernel_size][temp_index+offset[i]]), sizeof(Fiber));
								if (InnerProduct(base_dir, f[i+9].seed.dir) < 0.0f) {
									ReverseFiber(&(f[i+9]));
								}
								memcpy(&(f[i+18]), &(parameters[2+cz+kernel_size][temp_index+offset[i]]), sizeof(Fiber));
								if (InnerProduct(base_dir, f[i+18].seed.dir) < 0.0f) {
									ReverseFiber(&(f[i+18]));
								}
							}

							float temp_tensor[7] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
							ComputeTensor(f, 27, temp_tensor);

							int kx = kernel_size + cx;
							int ky = kernel_size + cy;
							int kz = kernel_size + cz;
							int gk_index = (kz * ks + ky) * ks + kx;

							st[(z*w*h+index)*7+0] += gaussian[gk_index] * temp_tensor[0];
							st[(z*w*h+index)*7+1] += gaussian[gk_index] * temp_tensor[1];
							st[(z*w*h+index)*7+2] += gaussian[gk_index] * temp_tensor[2];
							st[(z*w*h+index)*7+3] += gaussian[gk_index] * temp_tensor[3];
							st[(z*w*h+index)*7+4] += gaussian[gk_index] * temp_tensor[4];
							st[(z*w*h+index)*7+5] += gaussian[gk_index] * temp_tensor[5];
							st[(z*w*h+index)*7+6] += gaussian[gk_index] * temp_tensor[6];
						}
					}
				}
			}

			float ff = (float)(y - kernel_size - 1.0f) / float(h - 2.0f * kernel_size - 2.0f) * 100.0f;
			printf("\rz = %d  %.2f%% finish", z, ff);

		}

		for (int i = 0; i < 2*(kernel_size+1); ++i) {
			memcpy(parameters[i], parameters[i+1], sizeof(Fiber)*w*h);
		}
		memset(parameters[2*(kernel_size+1)], 0, sizeof(Fiber)*w*h);

		printf("\rz = %d  %.2f%% finish\n", z, 100.0f);
	}

	/* save the result */
	CDTI *pST = new CDTI();
	pST->CreateDTI(w, h, d);
	for (int z = 0; z < d; ++z) {
		for (int y = 0; y < h; ++y) {
			for (int x = 0; x < w; ++x) {
				int index = (z * h + y) * w + x;
				pST->SetDTI(x, y, z, &(st[index*7]));
			}
		}
	}
	pST->SaveDTIFile(dst_pathname);
	SafeDelete(pST);


	/* free memory */
	for (int z = 0; z < d; ++z) {
		SafeDeleteArray(confPathname[z]);
		SafeDeleteArray(vertPathname[z]);
	}
	SafeDeleteArray(confPathname);
	SafeDeleteArray(vertPathname);

	SafeDeleteArray(st);
	SafeDeleteArray(gaussian);

	for (int i = 0; i < 2*(kernel_size+1)+1; ++i) {
		SafeDeleteArray(parameters[i]);
	}
	SafeDeleteArray(parameters);

	return 0;
}